#!/usr/bin/bash

red_head="\e[31m"
red_end="\e[0m"

green_head="\e[32m"
green_end="\e[0m"

yellow_head="\e[33m"
yellow_end="\e[0m"

blue_head="\e[36m"
blue_end="\e[0m"

install_list=(cuda-toolkit cudnn tensorrt)
to_install=()


install_target()
{
	case $1 in 
		(cuda-toolkit)
			wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
			sudo mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
			wget https://developer.download.nvidia.com/compute/cuda/12.4.1/local_installers/cuda-repo-ubuntu2004-12-4-local_12.4.1-550.54.15-1_amd64.deb
			sudo dpkg -i cuda-repo-ubuntu2004-12-4-local_12.4.1-550.54.15-1_amd64.deb
			sudo cp /var/cuda-repo-ubuntu2004-12-4-local/cuda-*-keyring.gpg /usr/share/keyrings/
			sudo apt update
			sudo apt install -y cuda
			echo "export PATH=/usr/local/cuda-12.4/bin/:$PATH" >> ~/.bashrc
			echo "export LD_LIBRARY_PATH=/usr/local/cuda-12.4/lib64/:$LD_LIBRARY_PATH" >> ~/.bashrc
		;;
		(cudnn)
			wget https://developer.download.nvidia.com/compute/cudnn/9.4.0/local_installers/cudnn-local-repo-ubuntu2004-9.4.0_1.0-1_amd64.deb
			sudo dpkg -i cudnn-local-repo-ubuntu2004-9.4.0_1.0-1_amd64.deb
			sudo cp /var/cudnn-local-repo-ubuntu2004-9.4.0/cudnn-*-keyring.gpg /usr/share/keyrings/
			sudo apt update
			sudo apt install -y cudnn
		;;
		(tensorrt)
			sudo dpkg -i nv-tensorrt-local-repo-ubuntu2004-10.5.0-cuda-12.6_1.0-1_amd64.deb
			sudo cp /var/nv-tensorrt-local-repo-ubuntu2004-10.5.0-cuda-12.6/*-keyring.gpg /usr/share/keyrings/
			sudo apt update
			sudo apt install tensorrt
			;;
	esac
}

echo -e "${blue_head}Now, we are going to install the nvidia tools${blue_end}"
for m in ${install_list[*]}; do
	read -p "Do you want to install ${m} ? (yes/no)" condition
	if [ ${condition} == "yes" ]; then
		if [ ${m} == "tensorrt" ]; then
			echo -e "${yellow_head}Please ensure that you have downloaded the tensorrt local installation .deb package and place it with this file${yellow_end}"
			echo -e "nv-tensorrt-local-repo-ubuntu2004-10.5.0-cuda-12.6_1.0-1_amd64.deb"
			if [ -e nv-tensorrt-local-repo-ubuntu2004-10.5.0-cuda-12.6_1.0-1_amd64.deb ]; then
				to_install+=(${m})
			else
				echo -e "${red_head}Lack of the .deb file, please download it from https://developer.nvidia.com/tensorrt/download/10x${red_end}"
			fi
		else
			to_install+=(${m})
		fi
	fi
done

echo -e "${yellow_head}Now, we are going to install${yellow_end} ${blue_head}${to_install[*]}${blue_end}"
read -p "Continue ? (yes/no)" condition
if [ ${condition} == "yes" ]; then
	echo -e "Start installation"
	for m in ${to_install[*]}; do
		install_target ${m}
	done
fi
